"""
增强版 RabbitMQ 消息路由系统
支持智能体间直接通信、消息持久化、优先级队列、消息确认等高级特性
集成AgentPrune框架以降低冗余通信
"""
import asyncio
import json
import uuid
from typing import Dict, Any, Optional, Callable, List
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
from enum import Enum
import aio_pika
from aio_pika import Message, DeliveryMode, ExchangeType, RobustConnection
from utils.logger import Logger
from communication.fipa_acl import FIPAACLProtocol, ACLMessage, Performative
from communication.agent_registry import agent_registry, AgentInfo
from communication.agent_prune import agent_prune_manager

logger = Logger.get_logger("rabbitmq_enhanced")


class MessagePriority(Enum):
    """消息优先级"""
    LOW = 1
    NORMAL = 5
    HIGH = 10
    CRITICAL = 15


@dataclass
class EnhancedMessage:
    """增强消息结构"""
    message_id: str
    sender_id: str
    recipient_id: str
    message_type: str
    payload: Dict[str, Any]
    priority: MessagePriority
    ttl_seconds: int
    correlation_id: str
    reply_to: Optional[str] = None
    timestamp: datetime = None
    
    def __post_init__(self):
        if self.timestamp is None:
            self.timestamp = datetime.utcnow()
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "message_id": self.message_id,
            "sender_id": self.sender_id,
            "recipient_id": self.recipient_id,
            "message_type": self.message_type,
            "payload": self.payload,
            "priority": self.priority.value,
            "ttl_seconds": self.ttl_seconds,
            "correlation_id": self.correlation_id,
            "reply_to": self.reply_to,
            "timestamp": self.timestamp.isoformat()
        }


class EnhancedRabbitMQManager:
    """增强版 RabbitMQ 管理器"""
    
    def __init__(self, connection_url: str = "amqp://guest:guest@localhost/"):
        self.connection_url = connection_url
        self.connection: Optional[RobustConnection] = None
        self.channel = None
        self.exchanges = {}
        self.queues = {}
        self.message_handlers = {}
        self.fipa_protocol = FIPAACLProtocol()
        
        # 交换器配置
        self.exchange_configs = {
            "agent_direct": {"type": ExchangeType.DIRECT, "durable": True},
            "agent_broadcast": {"type": ExchangeType.FANOUT, "durable": True},
            "agent_topic": {"type": ExchangeType.TOPIC, "durable": True},
            "agent_priority": {"type": ExchangeType.HEADERS, "durable": True}
        }
        
        # 队列配置
        self.queue_configs = {
            "agent_messages": {
                "durable": True,
                "arguments": {
                    "x-message-ttl": 3600000,  # 1小时TTL
                    "x-max-priority": 15,      # 最大优先级
                    "x-dead-letter-exchange": "dlx_agent_messages",
                    "x-dead-letter-routing-key": "dead_letters"
                }
            },
            "fipa_messages": {
                "durable": True,
                "arguments": {
                    "x-message-ttl": 1800000,  # 30分钟TTL
                    "x-max-priority": 10
                }
            }
        }
    
    async def connect(self) -> bool:
        """建立连接"""
        try:
            self.connection = await aio_pika.connect_robust(self.connection_url)
            self.channel = await self.connection.channel()
            
            # 设置QoS，限制未确认消息数量
            await self.channel.set_qos(prefetch_count=100)
            
            # 声明交换器
            for exchange_name, config in self.exchange_configs.items():
                exchange = await self.channel.declare_exchange(
                    name=exchange_name,
                    type=config["type"],
                    durable=config["durable"]
                )
                self.exchanges[exchange_name] = exchange
            
            # 声明队列
            for queue_name, config in self.queue_configs.items():
                queue = await self.channel.declare_queue(
                    name=queue_name,
                    durable=config["durable"],
                    arguments=config.get("arguments", {})
                )
                self.queues[queue_name] = queue
            
            # 声明死信交换器和队列
            dlx_exchange = await self.channel.declare_exchange(
                name="dlx_agent_messages",
                type=ExchangeType.DIRECT,
                durable=True
            )
            
            dlx_queue = await self.channel.declare_queue(
                name="dead_letter_queue",
                durable=True
            )
            
            await dlx_queue.bind(dlx_exchange, "dead_letters")
            
            logger.info("增强版 RabbitMQ 连接成功")
            return True
            
        except Exception as e:
            logger.error(f"RabbitMQ 连接失败: {str(e)}")
            return False
    
    async def disconnect(self):
        """断开连接"""
        try:
            if self.connection:
                await self.connection.close()
                logger.info("RabbitMQ 连接已关闭")
        except Exception as e:
            logger.error(f"RabbitMQ 断开连接失败: {str(e)}")
    
    async def send_direct_message(self, sender_id: str, recipient_id: str, 
                                message_type: str, payload: Dict[str, Any],
                                priority: MessagePriority = MessagePriority.NORMAL,
                                ttl_seconds: int = 3600) -> Optional[str]:
        """发送直接消息"""
        try:
            message_id = str(uuid.uuid4())
            correlation_id = str(uuid.uuid4())
            
            # 高优先级消息不进行剪枝
            should_process = True
            prune_reason = None
            
            if priority.value < 8:  # 优先级低于8的消息进行剪枝判断
                # 使用AgentPrune进行消息剪枝判断
                should_process, prune_reason = await agent_prune_manager.process_message(
                    payload, sender_id, recipient_id
                )
                
                if not should_process:
                    logger.debug(f"消息被AgentPrune剪枝: message_id={message_id}, reason={prune_reason}")
                    # 返回消息ID，但实际上不会发送
                    return message_id
            
            enhanced_msg = EnhancedMessage(
                message_id=message_id,
                sender_id=sender_id,
                recipient_id=recipient_id,
                message_type=message_type,
                payload=payload,
                priority=priority,
                ttl_seconds=ttl_seconds,
                correlation_id=correlation_id
            )
            
            message = Message(
                json.dumps(enhanced_msg.to_dict()).encode(),
                message_id=message_id,
                correlation_id=correlation_id,
                priority=priority.value,
                expiration=timedelta(seconds=ttl_seconds),
                delivery_mode=DeliveryMode.PERSISTENT,
                headers={
                    "agent_prune_processed": not should_process,
                    "agent_prune_reason": prune_reason
                }
            )
            
            # 根据消息类型选择交换器和路由键
            if message_type == "fipa_acl":
                routing_key = f"fipa.{sender_id}.{recipient_id}"
                exchange = self.exchanges["agent_direct"]
            else:
                routing_key = f"direct.{sender_id}.{recipient_id}"
                exchange = self.exchanges["agent_direct"]
            
            await exchange.publish(message, routing_key=routing_key)
            
            logger.info(f"直接消息发送成功: {sender_id} -> {recipient_id}, type={message_type}, priority={priority.name}")
            return message_id
            
        except Exception as e:
            logger.error(f"直接消息发送失败: {sender_id} -> {recipient_id}, error={str(e)}")
            return None
    
    async def send_fipa_message(self, sender_id: str, recipient_id: str, 
                              performative: Performative, content: Dict[str, Any],
                              priority: MessagePriority = MessagePriority.NORMAL) -> Optional[str]:
        """发送 FIPA-ACL 消息"""
        try:
            # 创建 FIPA-ACL 消息
            acl_message = ACLMessage(
                sender=sender_id,
                recipients=[recipient_id],
                performative=performative,
                content=content,
                reply_with=str(uuid.uuid4()),
                language="JSON",
                ontology="maess_domain"
            )
            
            # 转换为增强消息格式
            payload = {
                "acl_message": acl_message.to_dict(),
                "protocol": "fipa_acl"
            }
            
            return await self.send_direct_message(
                sender_id=sender_id,
                recipient_id=recipient_id,
                message_type="fipa_acl",
                payload=payload,
                priority=priority,
                ttl_seconds=1800  # 30分钟TTL
            )
            
        except Exception as e:
            logger.error(f"FIPA-ACL 消息发送失败: {sender_id} -> {recipient_id}, error={str(e)}")
            return None
    
    async def broadcast_message(self, sender_id: str, message_type: str, 
                              payload: Dict[str, Any],
                              priority: MessagePriority = MessagePriority.NORMAL) -> Optional[str]:
        """广播消息"""
        try:
            message_id = str(uuid.uuid4())
            correlation_id = str(uuid.uuid4())
            
            enhanced_msg = EnhancedMessage(
                message_id=message_id,
                sender_id=sender_id,
                recipient_id="*",  # 广播给所有智能体
                message_type=message_type,
                payload=payload,
                priority=priority,
                ttl_seconds=1800,
                correlation_id=correlation_id
            )
            
            message = Message(
                json.dumps(enhanced_msg.to_dict()).encode(),
                message_id=message_id,
                correlation_id=correlation_id,
                priority=priority.value,
                delivery_mode=DeliveryMode.PERSISTENT
            )
            
            exchange = self.exchanges["agent_broadcast"]
            await exchange.publish(message, routing_key="broadcast")
            
            logger.info(f"广播消息发送成功: {sender_id}, type={message_type}, priority={priority.name}")
            return message_id
            
        except Exception as e:
            logger.error(f"广播消息发送失败: {sender_id}, error={str(e)}")
            return None
    
    async def register_message_handler(self, agent_id: str, message_type: str, 
                                     handler: Callable[[EnhancedMessage], None]):
        """注册消息处理器"""
        handler_key = f"{agent_id}:{message_type}"
        self.message_handlers[handler_key] = handler
        
        logger.info(f"消息处理器注册成功: {handler_key}")
    
    async def start_consuming(self, agent_id: str):
        """开始消费消息"""
        try:
            # 向AgentPrune注册智能体角色
            try:
                agent_info = agent_registry.get_agent_info(agent_id)
                if agent_info:
                    agent_prune_manager.register_agent_role(agent_id, agent_info.agent_type)
            except Exception as e:
                logger.warning(f"向AgentPrune注册智能体角色失败: {str(e)}")
                
            # 绑定队列到交换器
            await self.queues["agent_messages"].bind(
                self.exchanges["agent_direct"], 
                routing_key=f"direct.*.{agent_id}"
            )
            await self.queues["agent_messages"].bind(
                self.exchanges["agent_direct"], 
                routing_key=f"fipa.*.{agent_id}"
            )
            await self.queues["agent_messages"].bind(
                self.exchanges["agent_broadcast"], 
                routing_key="broadcast"
            )
            
            async def process_message(message: aio_pika.IncomingMessage):
                async with message.process():
                    try:
                        # 解析消息
                        msg_data = json.loads(message.body.decode())
                        enhanced_msg = EnhancedMessage(
                            message_id=msg_data["message_id"],
                            sender_id=msg_data["sender_id"],
                            recipient_id=msg_data["recipient_id"],
                            message_type=msg_data["message_type"],
                            payload=msg_data["payload"],
                            priority=MessagePriority(msg_data["priority"]),
                            ttl_seconds=msg_data["ttl_seconds"],
                            correlation_id=msg_data["correlation_id"],
                            reply_to=msg_data.get("reply_to"),
                            timestamp=datetime.fromisoformat(msg_data["timestamp"])
                        )
                        
                        # 调用处理器
                        handler_key = f"{agent_id}:{enhanced_msg.message_type}"
                        handler = self.message_handlers.get(handler_key)
                        
                        if handler:
                            await handler(enhanced_msg)
                        else:
                            # 默认处理器
                            await self._default_message_handler(enhanced_msg)
                            
                    except Exception as e:
                        logger.error(f"消息处理失败: {str(e)}")
            
            # 开始消费
            await self.queues["agent_messages"].consume(process_message)
            logger.info(f"消息消费启动成功: {agent_id}")
            
        except Exception as e:
            logger.error(f"消息消费启动失败: {agent_id}, error={str(e)}")
    
    async def _default_message_handler(self, message: EnhancedMessage):
        """默认消息处理器"""
        logger.info(f"默认处理消息: {message.message_id} from {message.sender_id} type={message.message_type}")
        
        # 如果是 FIPA-ACL 消息，使用 FIPA 协议处理
        if message.message_type == "fipa_acl" and "acl_message" in message.payload:
            acl_msg_data = message.payload["acl_message"]
            acl_message = ACLMessage.from_dict(acl_msg_data)
            
            # 处理 FIPA 消息
            response = await self.fipa_protocol.handle_message(acl_message)
            
            # 如果有回复，发送回去
            if response and message.reply_to:
                await self.send_direct_message(
                    sender_id=message.recipient_id,
                    recipient_id=message.sender_id,
                    message_type="fipa_acl",
                    payload={"acl_message": response.to_dict()},
                    correlation_id=message.correlation_id
                )
    
    async def request_reply(self, sender_id: str, recipient_id: str, 
                          message_type: str, payload: Dict[str, Any],
                          timeout_seconds: int = 30,
                          priority: MessagePriority = MessagePriority.NORMAL) -> Optional[EnhancedMessage]:
        """请求-回复模式"""
        try:
            # 创建临时回复队列
            reply_queue_name = f"reply_{sender_id}_{str(uuid.uuid4())[:8]}"
            reply_queue = await self.channel.declare_queue(
                name=reply_queue_name,
                exclusive=True,
                auto_delete=True
            )
            
            # 发送请求消息
            message_id = await self.send_direct_message(
                sender_id=sender_id,
                recipient_id=recipient_id,
                message_type=message_type,
                payload=payload,
                priority=priority,
                ttl_seconds=timeout_seconds
            )
            
            if not message_id:
                return None
            
            # 等待回复
            response_future = asyncio.Future()
            
            async def on_response(message: aio_pika.IncomingMessage):
                async with message.process():
                    try:
                        msg_data = json.loads(message.body.decode())
                        enhanced_msg = EnhancedMessage(
                            message_id=msg_data["message_id"],
                            sender_id=msg_data["sender_id"],
                            recipient_id=msg_data["recipient_id"],
                            message_type=msg_data["message_type"],
                            payload=msg_data["payload"],
                            priority=MessagePriority(msg_data["priority"]),
                            ttl_seconds=msg_data["ttl_seconds"],
                            correlation_id=msg_data["correlation_id"],
                            reply_to=msg_data.get("reply_to"),
                            timestamp=datetime.fromisoformat(msg_data["timestamp"])
                        )
                        response_future.set_result(enhanced_msg)
                    except Exception as e:
                        response_future.set_exception(e)
            
            await reply_queue.consume(on_response)
            
            # 等待回复或超时
            try:
                response = await asyncio.wait_for(response_future, timeout=timeout_seconds)
                return response
            except asyncio.TimeoutError:
                logger.warning(f"请求超时: {sender_id} -> {recipient_id}")
                return None
                
        except Exception as e:
            logger.error(f"请求-回复失败: {sender_id} -> {recipient_id}, error={str(e)}")
            return None
    
    async def get_queue_stats(self) -> Dict[str, Any]:
        """获取队列统计信息"""
        try:
            stats = {}
            
            for queue_name, queue in self.queues.items():
                queue_info = await queue.get_info()
                stats[queue_name] = {
                    "messages": queue_info.messages,
                    "consumers": queue_info.consumers,
                    "memory": queue_info.memory
                }
            
            # 添加AgentPrune的统计信息
            prune_stats = agent_prune_manager.get_pruning_stats()
            stats["agent_prune_stats"] = prune_stats
            
            return stats
            
        except Exception as e:
            logger.error(f"获取队列统计失败: {str(e)}")
            return {}


# 全局 RabbitMQ 管理器实例
rabbitmq_manager = EnhancedRabbitMQManager()


# 使用示例
async def example_usage():
    """增强版 RabbitMQ 使用示例"""
    
    # 连接到 RabbitMQ
    connected = await rabbitmq_manager.connect()
    if not connected:
        print("连接失败")
        return
    
    # 注册消息处理器
    async def sample_handler(message: EnhancedMessage):
        print(f"收到消息: {message.message_id} from {message.sender_id}")
        print(f"内容: {message.payload}")
        
        # 发送回复
        if message.reply_to:
            await rabbitmq_manager.send_direct_message(
                sender_id=message.recipient_id,
                recipient_id=message.sender_id,
                message_type="response",
                payload={"status": "processed", "original_id": message.message_id}
            )
    
    await rabbitmq_manager.register_message_handler("agent_001", "test_message", sample_handler)
    
    # 开始消费消息
    await rabbitmq_manager.start_consuming("agent_001")
    
    # 发送测试消息
    message_id = await rabbitmq_manager.send_direct_message(
        sender_id="agent_002",
        recipient_id="agent_001",
        message_type="test_message",
        payload={"test": "Hello World!"},
        priority=MessagePriority.HIGH
    )
    
    print(f"发送消息 ID: {message_id}")
    
    # 获取队列统计
    stats = await rabbitmq_manager.get_queue_stats()
    print(f"队列统计: {stats}")
    
    # 等待一段时间
    await asyncio.sleep(5)
    
    # 断开连接
    await rabbitmq_manager.disconnect()


if __name__ == "__main__":
    asyncio.run(example_usage())